import cv2 as cv
import numpy as np
from PIL import Image
import os
import json

# 存储目录，人脸数据模型，人脸图片库，映射文件
# 人脸图像数据库、训练模型和名字映射文件路径
dataset_path = 'dataset'  # 人脸图像数据存储路径
trainer_path = 'trainer'  # 训练好的模型存储路径
names_mapping_path = 'names_mapping.json'  # 名字到 ID 的映射文件路径{'1': 'zhangsan', '2': 'lisi'}


def input_names():
    """
    输入名字并自动分配 ID，保存到名字映射文件中
    """
    # 读取已存在的名字到 ID 映射
    if os.path.exists(names_mapping_path):
        with open(names_mapping_path, 'r') as f:
            names_mapping = json.load(f)
    else:
        names_mapping = {}

    # 自动分配 ID
    new_id = max(names_mapping.values(), default=0) + 1
    name = input(f"请输入名字 (ID 将自动分配为 {new_id}): ")
    names_mapping[name] = new_id

    # 保存名字到 ID 映射到文件
    with open(names_mapping_path, 'w') as f:
        json.dump(names_mapping, f)
    return new_id


def main():
    # 如果数据集目录不存在，则创建
    if not os.path.exists(dataset_path):
        os.makedirs(dataset_path)
    if not os.path.exists(trainer_path):
        os.makedirs(trainer_path)
    # 后续可以增加前端页面，现在使用控制台操作进行替代
    while True:
        print("请选择一个操作：")
        print("1. 数据采集")
        print("2. 模型训练")
        print("3. 实时识别")
        print("4. 退出")
        choice = input("输入你的选择: ")
        if choice == '1':
            face_id = input_names()  # 输入名字并获取 ID
            # TODO 功能2,3后续编写
        elif choice == '4':
            print("退出程序")
            break
        else:
            print("无效的选择，请重试。")


if __name__ == '__main__':
    main()
